import pprint
import logzero
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import font_manager
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure
from matplotlib.ticker import EngFormatter, FormatStrFormatter
from PyQt5.QtWidgets import QApplication, QSizePolicy, QMainWindow, QWidget, QVBoxLayout
from PyQt5.QtCore import Qt

FONTP = font_manager.FontProperties()
FONTP.set_family("SimHei")
FONTP.set_size(14)
YLABELS = (
    "A相电流",
    "B相电流",
    "C相电流",
    "UAB线电压",
    "UBC线电压",
    "UCA线电压",
    "正序电流",
    "负序电流",
    "开出1",
    "开出2",
    "开出3",
    "开出4",
    "开出5",
    "开出6",
    "开出7",
    "开出8",
)
YUNITS = ("A", "A", "A", "V", "V", "V", "A", "A", "", "", "", "", "", "", "", "")


# FigureCanvas inherits QWidget
class GraphCanvas(FigureCanvas):
    """
    GraphCanvas Class
    change window title should after QWidget.show
    """

    def __init__(
        self,
        titleHead="",
        figureTitile="",
        xlabel="",
        ylabel="",
        xlabelUnit="",
        xdata=np.arange(0, 4 * np.pi, 0.1),
        ydatas_0=[np.sin(np.arange(0, 4 * np.pi, 0.1))] * 8,
        ydatas_1=[np.sin(np.arange(0, 4 * np.pi, 0.1))] * 8,
        logger=logzero.logger,
        parent=None,
    ):
        self.logger = logger
        self.papa = parent
        self.logger.info("papa is {}".format(self.papa))
        # 显示样式
        # plt.style.use(["ggplot", "dark_background"])
        plt.style.use(["bmh"])
        self.fig = Figure(figsize=(8.9, 6.4), dpi=100)
        self.axes = self.fig.add_subplot(111)
        # 数据
        self.x = xdata
        self.ydatas_0 = ydatas_0
        self.ydatas_1 = ydatas_1
        self.y = self.ydatas_0[0]
        self.dataSetIndex = 0
        self.dataSetLength = len(ydatas_0) + len(ydatas_1)
        # 标签 标题
        self.fig.suptitle(figureTitile, fontproperties=FONTP)
        # self.axes.set_title(axesTitle, fontproperties=FONTP)
        self.axes.set_xlabel(xlabel, fontproperties=FONTP)
        self.axes.xaxis.set_major_formatter(EngFormatter(unit=xlabelUnit))
        self.axes.xaxis.set_minor_formatter(EngFormatter(unit=xlabelUnit))
        self.axes.set_ylabel(YLABELS[self.dataSetIndex], fontproperties=FONTP)
        self.axes.yaxis.set_major_formatter(EngFormatter(unit=YUNITS[self.dataSetIndex]))
        # 连续线
        self.line = self.axes.plot(self.x, self.y, color="g", linewidth=1.0)[0]
        # 散点图
        self.c = np.random.randint(1, 5, size=self.x.size)
        self.norm = plt.Normalize(1, 4)
        self.cmap = plt.cm.RdYlGn
        self.sc = self.axes.scatter(self.x, self.y, c="r", s=10, cmap=self.cmap, norm=self.norm)
        self.st = None
        # 鼠标点击记录
        self.annotKeepDict = {}
        self.annotLastClickList = []
        # 按键
        self.pressIdx = -1
        self.annotLastKeyPress = None

        super(GraphCanvas, self).__init__(self.fig)
        self.setParent(parent)
        # 窗口标题
        self.titleHead = titleHead
        self.updateAxesTitle()

        self.fig.canvas.mpl_connect("key_press_event", self.on_press)
        self.fig.canvas.mpl_connect("button_press_event", self.on_click)
        # https://stackoverflow.com/questions/11551049/matplotlib-plot-zooming-with-scroll-wheel
        self.fig.canvas.mpl_connect("scroll_event", self.on_scroll)
        self.fig.canvas.mpl_connect("motion_notify_event", self.on_hover)

        FigureCanvas.setSizePolicy(self, QSizePolicy.Expanding, QSizePolicy.Expanding)
        FigureCanvas.updateGeometry(self)

    def updateAxesTitle(self):
        newTitle = " {}  |  {} / {}".format(YLABELS[self.dataSetIndex], self.dataSetIndex + 1, self.dataSetLength)
        self.axes.set_title(newTitle, fontproperties=FONTP)

    def dealClickedKey(self, idx, updateLastRec=False):
        dc = self.annotKeepDict.get(idx)
        if dc is not None:
            if dc.get("clicked") is True:
                dc["clicked"] = False
                dc["annotate"].set_visible(False)
                if updateLastRec and idx in self.annotLastClickList:
                    self.annotLastClickList.remove(idx)
            else:
                dc["clicked"] = True
                dc["annotate"].set_visible(True)
                if updateLastRec:
                    self.annotLastClickList.append(idx)

    def dealClickedKeyAll(self, updateLastRec):
        for idx in self.annotKeepDict.keys():
            self.dealClickedKey(idx, updateLastRec)

    def on_click(self, event):
        # self.logger.info('%s click: button=%d, x=%d, y=%d, xdata=%f, ydata=%f' %
        #       ('double' if event.dblclick else 'single', event.button,
        #         event.x, event.y, event.xdata, event.ydata))
        # self.logger.info('on click')
        if event.dblclick and event.button == 2:
            self.logger.info("double click {}".format(pprint.pformat(self.annotKeepDict)))
            for dc in self.annotKeepDict.values():
                if dc.get("clicked") is True:
                    dc["clicked"] = False
                    dc["annotate"].set_visible(False)
            self.annotLastClickList = []
        elif not event.dblclick:
            if event.button == 3:
                self.logger.info("right click {}".format(self.annotLastClickList))
                if len(self.annotLastClickList) > 0:
                    idx = self.annotLastClickList.pop()
                    dc = self.annotKeepDict.get(idx)
                    dc["clicked"] = False
                    dc["annotate"].set_visible(False)
            elif event.button == 1:
                self.logger.info("left click {}".format(self.annotLastClickList))
                cont, ind = self.sc.contains(event)
                if cont:
                    idx = ind["ind"][0]
                    self.dealClickedKey(idx, updateLastRec=True)
                self.logger.info("cont {}, ind {}".format(cont, ind))
        self.fig.canvas.draw_idle()

    def on_scroll(self, event):
        self.logger.info("scroll {} step {}".format(event.button, event.step))
        newIndex = int(self.dataSetIndex - event.step)
        if newIndex >= self.dataSetLength:
            self.dataSetIndex = self.dataSetLength - 1
        elif newIndex < 0:
            self.dataSetIndex = 0
        else:
            self.dataSetIndex = newIndex
            self.dealDataSetIndexChange()

    def dealDataSetIndexChange(self):
        """https://stackoverflow.com/questions/4098131/how-to-update-a-plot-in-matplotlib"""
        for dc in self.annotKeepDict.values():
            annotate = dc["annotate"]
            annotate.set_visible(False)
        self.annotKeepDict.clear()
        self.annotLastClickList.clear()
        self.sc.remove()
        if self.st is not None and len(self.st) > 0:
            for l2 in self.st:
                l2.set_visible(False)
        if self.dataSetIndex < len(self.ydatas_0):
            self.y = self.ydatas_0[self.dataSetIndex]
            self.sc = self.axes.scatter(self.x, self.y, c="r", s=10, cmap=self.cmap, norm=self.norm)
            self.line.set_ydata(self.y)
            self.line.set_visible(True)
            self.axes.yaxis.set_major_formatter(EngFormatter(unit=YUNITS[self.dataSetIndex]))
        else:
            self.line.set_visible(False)
            self.y = self.ydatas_1[self.dataSetIndex - len(self.ydatas_0)]
            self.st = self.axes.step(self.x, self.y, where="post")
            self.sc = self.axes.scatter(self.x, self.y, c="r", s=10, cmap=self.cmap, norm=self.norm)
            self.axes.yaxis.set_major_formatter(FormatStrFormatter("%1.2f"))
        self.axes.set_ylabel(YLABELS[self.dataSetIndex], fontproperties=FONTP)
        self.axes.relim(visible_only=True)
        self.axes.autoscale(axis="y")
        self.fig.canvas.draw_idle()
        self.updateAxesTitle()

    def on_hover(self, event):
        # self.logger.info('on hover')
        if event.inaxes == self.axes and not isinstance(self.sc, list):
            cont, ind = self.sc.contains(event)
            # self.logger.info('cont {}, ind {}'.format(cont, ind))
            if cont:
                self.update_annot(ind)
                self.pressIdx = ind["ind"][0]
            for idx, dc in self.annotKeepDict.items():
                clicked = dc.get("clicked")
                annotate = dc["annotate"]
                if (clicked is True) or (cont and idx == ind["ind"][0]):
                    annotate.set_visible(True)
                else:
                    annotate.set_visible(False)

            self.fig.canvas.draw_idle()

    def update_annot(self, ind):
        # import pdb; pdb.set_trace()
        idx = ind["ind"][0]
        self.getOrNewAnnotate(idx)

    def getOrNewAnnotate(self, idx):
        dc = self.annotKeepDict.get(idx, {})
        if dc.get("annotate") is None:
            pos = self.sc.get_offsets()[idx]
            annot = self.axes.annotate(
                "",
                xy=pos,
                xytext=(-35, 20),
                textcoords="offset points",
                arrowprops=dict(arrowstyle="->"),
                bbox=dict(boxstyle="round4", fc="w"),
            )
            text = "x: {:.1f} y: {:.4f}".format(self.x[idx], self.y[idx])
            annot.set_text(text)
            annot.get_bbox_patch().set_facecolor(self.cmap(self.norm(self.c[idx])))
            annot.get_bbox_patch().set_alpha(0.4)
            dc["annotate"] = annot
            self.annotKeepDict[idx] = dc
        return dc

    def on_press(self, event):
        """
        deal key press
        """
        keyCode = event.key
        self.logger.info("press {}".format(keyCode))
        if keyCode in ("a", "left", "d", "right"):
            # modify self.pressIdx
            if keyCode in ("a", "left"):
                if self.pressIdx <= 0:
                    self.pressIdx = self.x.size - 1
                else:
                    self.pressIdx -= 1
            elif keyCode in ("d", "right"):
                if self.pressIdx >= self.x.size - 1:
                    self.pressIdx = 0
                else:
                    self.pressIdx += 1

            if isinstance(self.annotLastKeyPress, dict):
                if self.annotLastKeyPress.get("annotate") is not None:
                    if self.annotLastKeyPress.get("clicked") is True:
                        self.annotLastKeyPress["annotate"].set_visible(True)
                    else:
                        self.annotLastKeyPress["annotate"].set_visible(False)
            dc = self.getOrNewAnnotate(self.pressIdx)
            self.logger.info("key press dc is {}".format(dc))
            self.annotLastKeyPress = dc
            self.annotLastKeyPress["annotate"].set_visible(True)
            self.fig.canvas.draw_idle()
        elif keyCode in ("w", "up", "s", "down"):
            delta = 1
            if keyCode in ("w", "up"):
                delta = -1
            newIndex = self.dataSetIndex + delta
            if newIndex >= self.dataSetLength:
                self.dataSetIndex = self.dataSetLength - 1
            elif newIndex < 0:
                self.dataSetIndex = 0
            else:
                self.dataSetIndex = newIndex
                self.dealDataSetIndexChange()
        self.logger.info("self.pressIdx is {}".format(self.pressIdx))


class LineGraph(QWidget):
    def __init__(self, confParams, *args, **kwargs):
        QWidget.__init__(self, *args, **kwargs)
        self.setParent(kwargs.get("parent"))
        confParams["parent"] = self
        self.logger = confParams.get("logger", logzero.logger)
        self.setLayout(QVBoxLayout())
        self.canvas = GraphCanvas(**confParams)
        # https://stackoverflow.com/questions/31656345/matplotlib-key-press-event-does-not-respond
        # https://github.com/matplotlib/matplotlib/issues/707
        self.canvas.setFocusPolicy(Qt.ClickFocus)
        self.canvas.setFocus()
        self.toolbar = NavigationToolbar(self.canvas, self)
        self.layout().addWidget(self.toolbar)
        self.layout().addWidget(self.canvas)


class LineGraphMainWindow(QMainWindow):
    def __init__(self):
        QMainWindow.__init__(self)

        self.setAttribute(Qt.WA_DeleteOnClose)  # Garbage-collect the window after it's been closed.
        self.setWindowTitle("Test ([*]/9)")
        # self.setWindowModified(True)

        self.xdata = np.arange(0, 10 * np.pi, 0.1)
        ydatas_0 = (
            np.sin(self.xdata),
            np.cos(self.xdata),
            np.sin(self.xdata * 2),
            np.cos(self.xdata * 2),
            np.sin(self.xdata * 4),
            np.cos(self.xdata * 4),
            np.sin(self.xdata) / 2,
            np.cos(self.xdata / 2),
        )
        ydatas_1 = [self.testDataFunc(period=i + 2, sub=i) for i in range(8)]
        confParams = dict(
            titleHead="windowTitle",
            figureTitile="figureTitile",
            xdata=self.xdata,
            ydatas_0=ydatas_0,
            ydatas_1=ydatas_1,
            xlabel="时间",
            ylabel="ylabel",
            xlabelUnit="mS",
            parent=self,
        )
        main_widget = LineGraph(confParams=confParams)
        self.setCentralWidget(main_widget)
        self.show()

    def testDataFunc(self, period=2, sub=0):
        return np.array([1 if i % period <= sub else 0 for i, _ in enumerate(self.xdata)])


if __name__ == "__main__":
    import sys

    app = QApplication(sys.argv)
    app_window = LineGraphMainWindow()
    sys.exit(app.exec_())
